import os
from flask_restful import Resource
from flask import request
from common.resp import ajax_success, ajax_error, ajax_data, ajax_page
from core import db, app
from datetime import datetime

class DatasetModel(db.Model):
    __tablename__ = "sys_dataset"

    id = db.Column(db.Integer, autoincrement=True, primary_key=True)
    name = db.Column(db.String(50), nullable=False, index=True)
    code = db.Column(db.String(50), nullable=False, index=True)
    remark = db.Column(db.Text, nullable=True)
    path = db.Column(db.String(255), nullable=False)
    classes = db.Column(db.JSON, nullable=False)
    image_count = db.Column(db.Integer, default=0)
    status = db.Column(db.String(10), default="0")  # 0空, 1采集中, 2已采集, 3已标注, 4已训练
    create_time = db.Column(db.DateTime, default=datetime.now)
    update_time = db.Column(db.DateTime, default=datetime.now, onupdate=datetime.now)
    
    def to_json(self):
        return {
            'id': self.id,
            'name': self.name,
            'code': self.code,
            'remark': self.remark,
            'path': self.path,
            'classes': self.classes,
            'image_count': self.image_count,
            'status': self.status,
            'create_time': self.create_time.strftime('%Y-%m-%d %H:%M:%S'),
            'update_time': self.update_time.strftime('%Y-%m-%d %H:%M:%S'),
        }

class DatasetResource(Resource):
    def get(self):
        query = request.args;
        filters = []
        if query.get('name'):
            filters.append(DatasetModel.name.like('%' + query['name'] + '%'))
        if query.get('status'):
            filters.append(DatasetModel.status == query['status'])

        paginate = DatasetModel.query.filter(
            *filters).paginate(
            page=int(query.get('page', 1)), per_page=int(query.get('limit', 10)), error_out=False)

        return ajax_page([item.to_json() for item in paginate.items], paginate.total)

    def post(self):
        data = request.get_json()
        dataset = DatasetModel(
            name = data.get('name'),
            code = data.get('code'),
            remark = data.get('remark', ''),
            classes = data.get('classes', []),
            image_count = data.get('image_count', 0),
            status = data.get('status', '0'),
            path = data.get('path', data.get('code')),
        )
        dsPath = os.path.join(app.config['DATASET_PATH'], dataset.path)
        if not os.path.exists(dsPath):
            os.makedirs(dsPath)

        db.session.add(dataset)
        db.session.commit()
        return ajax_success('数据集创建成功')

    def put(self):
        data = request.get_json()
        if not data.get('id'):
            return ajax_error('参数错误')

        model = DatasetModel.query.get_or_404(data.get('id'))
        if not model:
            return ajax_error('数据集不存在')

        if data.get('name'):
            model.name = data.get('name', model.name)
        if data.get('remark'):
            model.remark = data.get('remark', model.remark)
        if data.get('status'):
            model.status = data.get('status', model.status)
        if data.get('classes'):
            model.classes = data.get('classes', model.classes)
        if data.get('path'):
            model.path = data.get('path', model.path)

        db.session.add(model)
        db.session.commit()
        return ajax_success('数据集更新成功')


class DatasetDetail(Resource):
    def get(self, id):
        if id:
            model = DatasetModel.query.get_or_404(id)
            return ajax_data(model.to_json())
        else:
            return ajax_error('参数错误')

    def delete(self, id):
        from modules.training import TrainingModel
        t = TrainingModel.query.filter_by(ds_id=id).first()
        if t:
            return ajax_error('删除数据集前请先删除训练任务')

        from modules.collection import CollectionModel
        t = CollectionModel.query.filter_by(ds_id=id).first()
        if t:
            return ajax_error('删除数据集前请先删除采集任务')

        # 先删除目录
        dataset = DatasetModel.query.get_or_404(id)
        dsPath = os.path.join(app.config['DATASET_PATH'], dataset.path)
        if os.path.exists(dsPath):
            os.rmdir(os.path.abspath(dsPath))

        db.session.delete(dataset)
        db.session.commit()
        return ajax_success('数据集删除成功')